package com.ruoyi.websocket;

import com.ruoyi.common.utils.ExceptionUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.websocket.Session;
import java.io.IOException;
import java.util.HashSet;
import java.util.Map;
import java.util.Objects;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;

/**
 * description: websocket管理器
 * 与 {@link WebSocketServer} 建立连接之后，session信息在该类中保存管理
 * auth fcb
 * date 2025-04-02 17:20
 **/
public class WebSocketManager {

    private static final Logger log = LoggerFactory.getLogger(WebSocketManager.class);

    private final static ConcurrentHashMap<String, Session> webSocketSessionMap = new ConcurrentHashMap<>();

    private final static ConcurrentHashMap<String, String> largeScreenAreaMap = new ConcurrentHashMap<>();

    public static void addSession(String sessionKey, Session session) {
        if(webSocketSessionMap.containsKey(sessionKey)) {
            log.warn("【websocket添加sessionKey重复】，sessionKey：{}", sessionKey);
            try {
                webSocketSessionMap.remove(sessionKey).close();
            } catch (IOException e) {
                log.error("【websocket关闭session异常】，sessionKey：{}, 异常：{}", sessionKey, ExceptionUtil.getExceptionMessage(e));
            }
        }
        webSocketSessionMap.put(sessionKey, session);
    }

    public static void setLargeScreenArea(String sessionKey, String area) {
        largeScreenAreaMap.put(sessionKey, area);
    }

    public static void removeLargeScreenArea(String sessionKey) {
        largeScreenAreaMap.remove(sessionKey);
    }

    public static void removeSession(String sessionKey) {
        webSocketSessionMap.remove(sessionKey);
    }

    public static Set<String> getSessionKeysByArea(String area) {
        Set<String> sessionKeys = new HashSet<>();
        for (Map.Entry<String, String> entry : largeScreenAreaMap.entrySet()) {
            if (Objects.equals(entry.getValue(), area)) {
                sessionKeys.add(entry.getKey());
            }
        }
        return sessionKeys;
    }

    public static Set<String> getSessionKeysByProvinceAndCity(String province, String city) {
        Set<String> sessionKeys = new HashSet<>();
        for (Map.Entry<String, String> entry : largeScreenAreaMap.entrySet()) {
            if (Objects.equals(entry.getValue(), province)) {
                sessionKeys.add(entry.getKey());
            } else if (Objects.equals(entry.getValue(), city)) {
                sessionKeys.add(entry.getKey());
            }
        }
        return sessionKeys;
    }

    public static Set<String> getOtherSessionKeys(Set<String> paramSessionKeys) {
        Set<String> sessionKeys = new HashSet<>(largeScreenAreaMap.keySet());
        sessionKeys.removeAll(paramSessionKeys);
        return sessionKeys;
    }

    public static Session getSession(String sessionKey) {
        return webSocketSessionMap.get(sessionKey);
    }

    /**
     * 发送消息给指定sessionKey
     * @param sessionKey websocket绘画标识，推荐使用 userId-businessType 的形式，例如：123-bigScreen
     * @param message 需要发送的数据
     */
    public static void sendMessageToKey(String sessionKey, String message) {
        Session session = webSocketSessionMap.get(sessionKey);
        if (session != null && session.isOpen()) {
            try {
                session.getBasicRemote().sendText(message);
            } catch (Exception e) {
                log.error("【websocket发送信息异常】，用户：{}, message：{}, 异常：{}", sessionKey, message, ExceptionUtil.getExceptionMessage(e));
            }
        }else {
            log.warn("【websocket发送信息未找到session】，用户：{}, message：{}", sessionKey, message);
        }
    }

    /**
     *  批量发送消息给所有sessionKey。
     * @param message 需要发送的数据
     */
    public static void sendMessageToAllKey(String message) {
        webSocketSessionMap.forEach((userId, session) -> {
            if (session != null && session.isOpen()) {
                try {
                    session.getBasicRemote().sendText(message);
                } catch (Exception e) {
                    log.error("【websocket批量发送信息异常】，用户：{}, message：{}, 异常：{}", userId, message, ExceptionUtil.getExceptionMessage(e));
                }
            }else {
                log.warn("【websocket批量发送信息未找到session】，用户：{}, message：{}", userId, message);
            }
        });
    }

    /**
     * 给指定用户发送消息。该方法只能用于sessionKey是 userId-businessType 的形式
     * @param userId 用户id
     * @param message 数据
     */
    public static void sendMessageToUserId(String userId, String message) {
        webSocketSessionMap.forEach((sessionKey, session) -> {
            String[] split = sessionKey.split("-");
            if (session != null && session.isOpen() && split[0].equals(userId)) {
                try {
                    session.getBasicRemote().sendText(message);
                } catch (Exception e) {
                    log.error("【websocket发送信息给指定用户异常】，用户：{}, message：{}, 异常：{}", userId, message, ExceptionUtil.getExceptionMessage(e));
                }
            }
        });
    }

    /**
     * 给指定业务发送消息。该方法只能用于sessionKey是 userId-businessType 的形式
     * @param businessType 业务标识
     * @param message 数据
     */
    public static void sendMessageToBusiness(String businessType, String message) {
        webSocketSessionMap.forEach((sessionKey, session) -> {
            String[] split = sessionKey.split("-");
            if (session != null && session.isOpen() && split[1].equals(businessType)) {
                try {
                    session.getBasicRemote().sendText(message);
                } catch (Exception e) {
                    log.error("【websocket发送信息给指定业务异常】，业务标识：{}, message：{}, 异常：{}", businessType, message, ExceptionUtil.getExceptionMessage(e));
                }
            }
        });
    }

    public static String getSessionKey(String userId, String businessType) {
        return userId + "-" + businessType;
    }
}
